﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Roslyn.Compilers.CSharp;
using System.Threading;
using Util;
using CompleterLib.Util;

namespace CompleterLib.AbstractTypes
{
    class ConstraintCollector : Roslyn.Compilers.CSharp.SyntaxWalker
    {
        private readonly Compilation comp;
        private static readonly Cache<Compilation, ConstraintCollector> ccs
            = new Cache<Compilation, ConstraintCollector>(
                comp => new ConstraintCollector(comp));
        private SemanticModel semModel;
        private List<TypeConstraint> constraints = new List<TypeConstraint>();

        private ConstraintCollector(Compilation comp)
        {
            this.comp = comp;
            foreach(var tree in comp.SyntaxTrees)
            {
                this.semModel = comp.GetSemanticModel(tree);
                var ctoken = new CancellationToken();
                this.Visit(tree.GetRoot(ctoken));
            }
        }
        public static ConstraintCollector GetForCompilation(Compilation comp)
        {
            return ccs[comp];
        }
        public IList<TypeConstraint> Constraints { get { return constraints.AsReadOnly(); } }
        private void AddConstraint(TypeConstraint c)
        {
            if (c != null)
            {
                constraints.Add(c);
            }
        }
        private void AddConstraints(IEnumerable<TypeConstraint> cs)
        {
            constraints.AddRange(cs.Where(c => c != null));
        }

        private Dictionary<ExpressionSyntax, TypeTerm> currentFuncExps
            = new Dictionary<ExpressionSyntax,TypeTerm>();
        internal TypeTerm GetTypeTermForDefinition(Symbol sym)
        {
            sym = sym.ResolveAlias();
            if (sym == null)
            {
                return null;
            }
            else if(sym is LocalSymbol)
            {
                return new NamedTypeTerm(sym as LocalSymbol);
            }
            else if (sym is FieldSymbol)
            {
                var field = sym as FieldSymbol;
                return field.AsTypeTerm();
            }
            else if(sym is PropertySymbol)
            {
                var prop = sym as PropertySymbol;
                return prop.AsTypeTerm();
            }
            else if (sym is ParameterSymbol)
            {
                var param = sym as ParameterSymbol;
                // TODO Index?
                // Isn't meth always currentMethod?
                var meth = param.ContainingSymbol as MethodSymbol;
                if (meth == null)
                {
                    throw new Exception("Unknown method for parameter");
                }
                return new FormalParameterTypeTerm(meth,
                    (meth.IsStatic ? 0 : 1) + param.Ordinal);
            }
            else if (sym is MethodSymbol)
            {
                if ((sym as MethodSymbol).MethodKind.Equals(MethodKind.Constructor))
                {
                    // Constructor calls are too general for their own type terms.
                    return new AnonymousTypeTerm();
                }
                else
                {
                    return new ReturnTypeTypeTerm(sym as MethodSymbol);
                }
            }
            else if (sym is TypeSymbol)
            {
                return null;
            }
            else
            {
                throw new NotImplementedException("Unknown target definition type: "
                    + sym.GetType());
            }
        }
        private TypeTerm GetTypeTermFor(ExpressionSyntax exp)
        {
            return GetTypeTermFor(exp, semModel);
        }
        internal TypeTerm GetTypeTermFor(ExpressionSyntax exp, SemanticModel semModel)
        {
            try
            {
                TypeTerm res;
                if (!currentFuncExps.TryGetValue(exp, out res))
                {
                    if (exp is IdentifierNameSyntax || exp is MemberAccessExpressionSyntax)
                    {
                        res = GetTypeTermForDefinition(semModel.GetSymbolInfo(exp).Symbol);
                    }
                    else if (exp is ThisExpressionSyntax)
                    {
                        res = new FormalParameterTypeTerm(currentMethod, 0);
                    }
                    else if (exp.Kind.Equals(SyntaxKind.AssignExpression))
                    {
                        res = GetTypeTermFor((exp as BinaryExpressionSyntax).Left, semModel);
                    }
                    else if (exp is InvocationExpressionSyntax)
                    {
                        var call = exp as InvocationExpressionSyntax;
                        var sym = semModel.GetSymbolInfo(call.Expression).Symbol as MethodSymbol;
                        if (sym == null || sym.MethodKind.Equals(MethodKind.Constructor))
                        {
                            res = new AnonymousTypeTerm();
                        }
                        else
                        {
                            res = new ReturnTypeTypeTerm(sym);
                        }
                    }
                    else if (exp is LiteralExpressionSyntax)
                    {
                        var ctc = exp as LiteralExpressionSyntax;
                        var constantValOpt = semModel.GetConstantValue(exp);
                        if (constantValOpt.HasValue)
                        {
                            var constantVal = constantValOpt.Value;
                            if (constantVal is string)
                            {
                                res = new ConstantStringTypeTerm(constantVal as string);
                            }
                            else
                            {
                                res = new ConstantTypeTerm(constantVal);
                            }
                        }
                    }
                    else if (exp is ConditionalExpressionSyntax)
                    {
                        var cond = exp as ConditionalExpressionSyntax;
                        res = GetTypeTermFor(cond.WhenTrue, semModel);
                    }
                    /*
                else if (exp is ICreateObjectInstance
                    || exp is ICreateDelegateInstance
                    || exp is IConversion
                    || exp is ICastIfPossible
                    || exp is ITypeOf
                    || exp is ICheckIfInstance
                    || exp is IAddressOf
                    || exp is IAddressDereference
                    || exp is IUnaryOperation
                    || exp is IDefaultValue
                    || exp is IPopValue
                    || exp is IBinaryOperation)
                {
                    res = new AnonymousTypeTerm();
                }
                else if (exp is ICreateArray)
                {
                    // TODO Could make arrays related to the type they contain...
                    res = new AnonymousTypeTerm();
                }*/
                    else if (exp is ObjectCreationExpressionSyntax
                        || exp is BinaryExpressionSyntax
                        || exp is ArrayCreationExpressionSyntax
                        || exp is DefaultExpressionSyntax
                        || exp is CastExpressionSyntax
                        || exp is TypeOfExpressionSyntax
                        || exp is PostfixUnaryExpressionSyntax
                        || exp is PrefixUnaryExpressionSyntax)
                    {
                        res = new AnonymousTypeTerm();
                    }
                    else
                    {
                        res = new AnonymousTypeTerm();
                        /*throw new NotImplementedException("Unknown expression type: "
                            + exp.GetType());*/
                    }
                    currentFuncExps[exp] = res;
                }
                return res;
            }
            catch (Exception e)
            {
                Logging.Logger.PrintLine("Exception in GetTypeTermFor: "
                    + e + "\n" + e.StackTrace);
                return new AnonymousTypeTerm();
            }
        }

        private MethodSymbol currentMethod;
        private PropertySymbol currentProperty;
        public override void VisitMethodDeclaration(MethodDeclarationSyntax node)
        {
            currentFuncExps = new Dictionary<ExpressionSyntax, TypeTerm>();
            currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
            base.VisitMethodDeclaration(node);            
        }
        public override void VisitPropertyDeclaration(PropertyDeclarationSyntax node)
        {
            currentFuncExps = new Dictionary<ExpressionSyntax, TypeTerm>();
            currentMethod = null;
            currentProperty = semModel.GetDeclaredSymbol(node) as PropertySymbol;
            base.VisitPropertyDeclaration(node);
        }
        public override void VisitAccessorDeclaration(AccessorDeclarationSyntax node)
        {
            currentFuncExps = new Dictionary<ExpressionSyntax, TypeTerm>();
            if (node.Kind.Equals(SyntaxKind.AddAccessorDeclaration) || node.Kind.Equals(SyntaxKind.RemoveAccessorDeclaration))
            {
                currentMethod = null;
            }
            else
            {
                try
                {
                    currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
                }
                catch (NullReferenceException)
                {
                    currentMethod = null;
                    Logging.Logger.PrintLine("VisitAccessorDeclaration: NullReferenceException"
                        + " in document " + semModel.Compilation.SyntaxTrees.First().FilePath
                        + " for node: " + node);
                }
            }
            base.VisitAccessorDeclaration(node);
        }
        public override void VisitConstructorDeclaration(ConstructorDeclarationSyntax node)
        {
            currentFuncExps = new Dictionary<ExpressionSyntax, TypeTerm>();
            currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
            base.VisitConstructorDeclaration(node);
        }
        public override void VisitFieldDeclaration(FieldDeclarationSyntax node)
        {
            currentFuncExps = new Dictionary<ExpressionSyntax, TypeTerm>();
            currentMethod = null;
            base.VisitFieldDeclaration(node);
        }

        public override void VisitInvocationExpression(InvocationExpressionSyntax call)
        {
            base.VisitInvocationExpression(call);

            var method = call.Expression;
            var methodType = semModel.GetTypeInfo(method).Type;
            var methodSymbol = semModel.GetSymbolInfo(method).Symbol;

            if (!(methodType is ErrorTypeSymbol) && methodSymbol != null)
                //&& methodType.TypeKind.Equals(TypeKind.Delegate))
            {
                var meth = methodSymbol as MethodSymbol;
                if (meth == null)
                {
                    return; // XXX ?
                }
                var args = call.ArgumentList.Arguments.Select(a => a.Expression);
                var methodAccess = method as MemberAccessExpressionSyntax;
                if (methodAccess != null)
                {
                    var parent = methodAccess.Expression;
                    if (!(parent is PredefinedTypeSyntax))
                    {
                        // TODO Should parent be included here?
                        var parentInfo = semModel.GetSymbolInfo(parent);
                        if (!(parentInfo.Symbol.ResolveAlias() is TypeSymbol))
                        {
                            args = parent.Cons(args);
                        }
                    }
                }
                else if (!meth.IsStatic)
                {
                    // null means this
                    args = EnumerableUtil.Cons(null, args);
                }
                AddConstraints(args.Select((arg, pos) =>
                    new FormalParameterTypeTerm(meth, pos)
                        .IsAssignedFrom(arg == null
                            ? new FormalParameterTypeTerm(currentMethod, 0)
                            : GetTypeTermFor(arg))));
            }
        }
        public override void VisitVariableDeclarator(VariableDeclaratorSyntax node)
        {
            base.VisitVariableDeclarator(node);
            var init = node.Initializer;
            if (init != null && !init.IsMissing)
            {
                var value = init.Value;
                var ident = node.Identifier;
                TypeTerm identTT = null;
                Symbol sym = semModel.GetDeclaredSymbol(node);
                if (sym is FieldSymbol)
                {
                    identTT = (sym as FieldSymbol).AsTypeTerm();
                }
                else if (currentMethod != null)
                {
                    identTT = new NamedTypeTerm(currentMethod, ident.ValueText);
                }

                if (identTT != null)
                {
                    AddConstraint(identTT.IsAssignedFrom(GetTypeTermFor(value)));
                }
            }
        }
        public override void VisitConditionalExpression(ConditionalExpressionSyntax node)
        {
            base.VisitConditionalExpression(node);
            AddConstraint(GetTypeTermFor(node.WhenTrue).IsAssignedFrom(GetTypeTermFor(node.WhenFalse)));
        }
        public override void VisitReturnStatement(ReturnStatementSyntax node)
        {
            base.VisitReturnStatement(node);
            var exp = node.Expression;
            if(exp != null && !exp.IsMissing)
            {
                if (currentMethod != null)
                {
                    AddConstraint(new ReturnTypeTypeTerm(currentMethod).IsAssignedFrom(GetTypeTermFor(exp)));
                }
                else if (currentProperty != null)
                {
                    AddConstraint(currentProperty.AsTypeTerm().IsAssignedFrom(GetTypeTermFor(exp)));
                }
            }
        }
        public override void VisitBinaryExpression(BinaryExpressionSyntax node)
        {
            base.VisitBinaryExpression(node);
            if (node.Kind.Equals(SyntaxKind.AssignExpression))
            {
                AddConstraint(GetTypeTermFor(node.Left).IsAssignedFrom(GetTypeTermFor(node.Right)));
            }
        }
    }

    class TypeConstraint
    {
        private TypeTerm left, right;
        public TypeConstraint(TypeTerm left, TypeTerm right)
        {
            this.left = left;
            this.right = right;
        }
        public TypeTerm Left { get { return left; } }
        public TypeTerm Right { get { return right; } }

        public override string ToString()
        {
            return Left + " = " + Right;
        }

        public void Replace(TypeTerm var, TypeTerm replacement)
        {
            left = left.Replace(var, replacement);
            right = right.Replace(var, replacement);
        }
    }
    static class TypeConstraintUtil
    {
        public static TypeConstraint IsAssignedFrom(this TypeTerm lvalue, TypeTerm rvalue)
        {
            if (lvalue == null || rvalue == null)
            {
                return null;
            }
            else
            {
                return new TypeConstraint(lvalue, rvalue);
            }
        }

        public static TypeTerm GetReturnTypeMeaning(
            this Dictionary<IdentifiedTypeTerm, TypeTerm> typeMeanings,
            MethodSymbol func)
        {
            var retTT = new ReturnTypeTypeTerm(func);
            TypeTerm res;
            if (typeMeanings.TryGetValue(retTT, out res))
            {
                return res;
            }
            else
            {
                return retTT;
            }
        }
        public static TypeTerm GetParameterTypeMeaning(
            this Dictionary<IdentifiedTypeTerm, TypeTerm> typeMeanings,
            MethodSymbol func,
            int position)
        {
            var paramTT = new FormalParameterTypeTerm(func, position);
            TypeTerm res;
            if (typeMeanings.TryGetValue(paramTT, out res))
            {
                return res;
            }
            else
            {
                return paramTT;
            }
        }
    }


}
